package com.chatplus.application.aiprocessor.listener.impl;

import cn.bugstack.chatglm.model.ChatCompletionResponse;
import cn.bugstack.chatglm.model.EventType;
import cn.hutool.extra.spring.SpringUtil;
import com.chatplus.application.aiprocessor.listener.StreamEventSourceListener;
import com.chatplus.application.aiprocessor.enumeration.MessageTypeEnum;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.enumeration.AiPlatformEnum;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.sse.EventSource;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.web.socket.WebSocketSession;

/**
 * ChatGLM监听器
 *
 * @author chj
 * @date 2024/1/4
 **/
public class ChatGLMEventSourceListener extends StreamEventSourceListener {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(ChatGLMEventSourceListener.class);
    String reply = "";

    public ChatGLMEventSourceListener(WebSocketSession session, AiPlatformEnum platform, String prompt) {
        super(session, platform, prompt);
    }

    @Override
    public void onEvent(@NotNull EventSource eventSource, String id, String type, @NotNull String data) {
        LOGGER.message("ChatGLM返回数据").context("data", data).context("type", type).info();
        ObjectMapper objectMapper = SpringUtil.getBean(ObjectMapper.class);
        // type 消息类型，add 增量，finish 结束，error 错误，interrupted 中断
        try {
            ChatCompletionResponse response = objectMapper.readValue(data, ChatCompletionResponse.class);
            String content = response.getData();
            if (EventType.error.getCode().equals(type)) {
                chatAiUtilHandler.replyErrorMessage("服务端发生异常，等待管理员修复!",false);
                return;
            }
            if (StringUtils.isNotEmpty(content)) {
                reply += content;
                chatAiUtilHandler.sendMessage(content);
            }
            if (EventType.finish.getCode().equals(type) ||"[DONE]".equals(data)) {
                ChatCompletionResponse.Meta meta = objectMapper.readValue(response.getMeta(), ChatCompletionResponse.Meta.class);
                ChatCompletionResponse.Usage usage = objectMapper.readValue(response.getMeta(), ChatCompletionResponse.Usage.class);

                LOGGER.message("回答完成").context("meta",meta).context("usage",usage).info();
                chatAiUtilHandler.saveChatHistory(reply, (long) usage.getTotal_tokens());
                chatAiUtilHandler.sendMessage(MessageTypeEnum.WS_END.getValue());
            }
        } catch (JsonProcessingException e) {
            chatAiUtilHandler.replyErrorMessage("服务端发生异常，等待管理员修复!",false);
            LOGGER.message("JSON 解析错误").exception(e).error();
        }
    }
}